import os, sys, math
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 4)))
sys.path.insert(0, parent_path)

from base64 import b64encode
import traceback
import numpy as np
from scipy.interpolate import interp1d
import socket, json, cv2, yaml, time, copy
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.visualize import get_color_map_list
from core.algorithm.deeplearning.paddle.detectors.seg.seg_utils import (
    visualize as seg_vis,
)


class FTPProcessor(BaseProcessor):
    """FTP检测后处理方案"""

    def __init__(
        self,
        keys,
        Custom_cls_map,
        Custom_target_cls,
        Custom_meter_cls_map,
        Custom_meter_target_cls,
    ) -> None:
        self.keys = keys
        if isinstance(Custom_target_cls, str):
            with open(Custom_target_cls, "r") as f:
                self.target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.target_cls = set(Custom_target_cls)
        if isinstance(Custom_cls_map, str):
            with open(Custom_cls_map, "r") as f:
                self.cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.cls_map = Custom_cls_map
        self.color_list = get_color_map_list(len(self.cls_map))

        if isinstance(Custom_meter_target_cls, str):
            with open(Custom_meter_target_cls, "r") as f:
                self.meter_target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.meter_target_cls = set(Custom_meter_target_cls)
        if isinstance(Custom_meter_cls_map, str):
            with open(Custom_meter_cls_map, "r") as f:
                self.meter_cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.meter_cls_map = Custom_meter_cls_map
        self.meter_color_list = get_color_map_list(len(self.meter_cls_map))

    def init_check(self):
        assert "det_meter" in self.keys
        assert "seg" in self.keys
        assert "crop_kpt" in self.keys
        super().init_check()

    def __call__(self, data: Dict) -> Dict:

        batch_targets = data[self.keys["in"]].copy()
        crop_targets = data[self.keys["crop_targets"]].copy()
        det_meter_result = data[self.keys["det_meter"]]
        seg_result = data[self.keys["seg"]] if "seg" in data else []
        kpt_result = data[self.keys["crop_kpt"]] if "crop_kpt" in data else []
        batch_objects = []
        batch_vis = []
        # 外部表盘检测(原图)
        for img_idx, bboxes in enumerate(det_meter_result["boxes"]):
            batch_img = batch_targets[img_idx]
            # crop_img = copy.deepcopy(crop_targets[img_idx]["imgs"][0])
            draw_thickness = min(batch_img.shape[:2]) // 320
            img = copy.deepcopy(batch_img)
            img_objects = []
            # 获取原图中检测到的表盘数据(切片内表盘数据)
            # for bbox in bboxes:
            for det_idx, bbox in enumerate(bboxes):
                cls_id, score, x1, y1, x2, y2 = bbox
                cls_id = int(cls_id)
                object = {
                    "cls_name": self.meter_cls_map[cls_id],
                    "box": [cls_id, score, x1, y1, x2, y2],
                    "ratio": [],
                }
                # 语义分割
                if "seg" in data:
                    seg_mask = seg_result[img_idx]["masks"][det_idx].astype(np.uint8)
                    # 获取补全轮廓后刻度曲线的二维点集和表盘轮廓弧长
                    points, have_target = connect_linestrip(seg_mask)
                else:
                    img_objects.append(object)
                    continue
                # 关键点检测
                if have_target:
                    keypoint = kpt_result[img_idx]["keypoint"][det_idx]
                    # 获取关键点坐标(原图)
                    kpt_start_point, kpt_end_point, pointer_tail, pointer_head = (
                        get_keypoint(keypoint)
                    )
                    # 从指针尾部向指针头部作一条射线，判断射线与曲线是否相交
                    if is_intersect(pointer_tail, pointer_head, points):
                        # 求指针射线与曲线轮廓交点
                        intersection = find_intersection(
                            points, pointer_head, pointer_tail
                        )
                        # 计算指针曲线弧长占比(保留四位小数)
                        ratio_arclen = find_arc_length_to_point(points, intersection)
                        ratio_arclen = round(ratio_arclen, 4)
                        object["ratio"] = ratio_arclen
                        img_objects.append(object)
                        # draw label
                        text = "{:.4f}".format(ratio_arclen)
                        # 给原图表盘画框
                        img = cv2.rectangle(
                            img,
                            (int(x1), int(y1)),
                            (int(x2), int(y2)),
                            color=self.meter_color_list[cls_id],
                            thickness=draw_thickness,
                        )
                        text_scale = max(0.5, img.shape[0] / 1600.0)
                        img = cv2.putText(
                            img,
                            text,
                            (int(x1), int(y1)),
                            cv2.FONT_ITALIC,
                            text_scale,
                            (0, 255, 255),
                            thickness=draw_thickness,
                        )
                        # test
                        cv2.imwrite(
                            "/home/smartgis/workspace/project/smart_server/data/output/pointer/result.jpg",
                            img,
                        )
                    else:
                        img_objects.append(object)
                        continue
                else:
                    img_objects.append(object)
                    continue
            batch_objects.append(img_objects)
            batch_vis.append(img)
        data[self.keys["out_vis"]] = img
        data[self.keys["out_obj"]] = batch_objects

        return data

    def close(self):
        self.status = 0


def get_keypoint(keypoint):
    """获取各个关键点坐标
    input:
        关键点: keypoint
    return:
        kpt_start_point: 表盘轮廓起点
        kpt_end_point: 表盘轮廓终点
        pointer_tail: 表针尾部
        pointer_head: 表针头部
    """
    kpt_start_point = keypoint[0][:2].astype(np.int32)
    kpt_end_point = keypoint[1][:2].astype(np.int32)
    pointer_tail = keypoint[2][:2].astype(np.int32)
    pointer_head = keypoint[3][:2].astype(np.int32)
    return kpt_start_point, kpt_end_point, pointer_tail, pointer_head


def connect_linestrip(label_mask):
    """将相近轮廓连接成一个轮廓,获取完整曲线轮廓的二维点集
    input:
        读取图像: image
    return:
        表盘曲线轮廓二维点集: points
        表盘曲线轮廓弧长: total_arclen
    """

    # 二值化图像
    _, binary = cv2.threshold(label_mask, 0, 255, cv2.THRESH_BINARY)
    # test
    cv2.imwrite(
        "/home/smartgis/workspace/project/smart_server/data/output/pointer/binary_org.jpg",
        binary,
    )
    # 连接相邻轮廓
    mask_rebuild = seg_rebuild(label_mask=binary)
    if mask_rebuild == []:
        points = []
        have_target = False
        return points, have_target
    # test
    cv2.imwrite(
        "/home/smartgis/workspace/project/smart_server/data/output/pointer/mask_rebuild.jpg",
        mask_rebuild,
    )
    # 检测轮廓
    contours, hierarchy = cv2.findContours(
        mask_rebuild, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    # 提取最大的轮廓（假设连接后的主要曲线是最大的轮廓）
    contour = max(contours, key=cv2.contourArea)
    # 将轮廓点转换为二维点集
    points = contour.squeeze()
    have_target = True
    return points, have_target


def seg_rebuild(label_mask, radius=75):
    """将实线(solid line)的语义分割的结果进行优化, 将位置接近的线段连接起来, 从而形成完整的实线线段
    input:
        label_mask: 语义分割掩码图, 目前只支持单分类的掩码图
        radius: 轮廓点扫描半径

    逻辑上即对掩码图 label_mask 进行轮廓提取, 然后两两逐个比对所有轮廓, 对每个轮廓逐轮廓点检查一定半径内是否有其他的轮廓
    如果有, 则找出哪个轮廓距离本轮廓最近,并且将两个轮廓之间最接近的两个轮廓点连接起来
    """
    # 连接细小的轮廓
    done = False
    while not done:
        # 轮廓提取
        contours, _ = cv2.findContours(
            label_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
        )
        if contours == ():
            label_mask = []
            return label_mask
        # 连通域提取
        _, ids_mask, stats, centres_xy = cv2.connectedComponentsWithStats(label_mask)

        # 得到连通域和轮廓之间的映射表
        contour2connect = {}
        connect2contour = {}
        for contour_idx, contour in enumerate(contours):
            cX = contour[0, 0, 0]
            cY = contour[0, 0, 1]
            connect_idx = ids_mask[cY, cX]  # 得到该轮廓属于第几个连通域
            contour2connect[contour_idx] = connect_idx
            connect2contour[connect_idx] = contour_idx

        # 逐轮廓开始比对, 找出位置相近的轮廓并连接
        for contour_idx, contour in enumerate(contours):

            if contour.shape[0] < 3:  # 轮廓小于3个点, 不呈多边形
                continue

            cX = contour[0, 0, 0]
            cY = contour[0, 0, 1]
            connect_idx = contour2connect[contour_idx]
            centre_xy = centres_xy[connect_idx]

            if contour.shape[0] > 30:  # 当轮廓点过多时, 精简轮廓点
                contour = cv2.approxPolyDP(contour, 2, True)

            # 循环每个轮廓点, 检测其附近有没有别的连通域
            near_connect_idx = None
            distance = np.inf
            connect_point = None
            for point_ in contour:
                point_ = point_[0]
                mask = np.zeros_like(label_mask, dtype=np.uint8)
                mask = cv2.circle(
                    mask,
                    point_.astype(np.int32),
                    radius=radius,
                    color=255,
                    thickness=-1,
                )
                check_list = set(ids_mask[mask > 0].tolist())
                for (
                    check_
                ) in check_list:  # 检查这些连通域距离当前轮廓的距离, 找出最近连通域
                    if check_ in [0, connect_idx]:  # 跳过背景类和自身 0
                        continue
                    distance_ = np.sqrt(
                        np.sum(np.power(centres_xy[check_] - centre_xy, 2))
                    )  # 计算两点距离
                    if distance_ < distance:
                        distance = distance_
                        near_connect_idx = check_
                        connect_point = point_.copy()

            if near_connect_idx is None:  # 附近没有别的轮廓了
                if contour_idx == len(contours) - 1:  # 已经执行完所有的轮廓, 结束
                    done = True
                    break
                continue  # 当前轮廓附近没有别的连通域

            # 附近存在别的轮廓, 根据连通域id取出该轮廓的数据
            near_contour_idx = connect2contour[near_connect_idx]
            near_contour = contours[near_contour_idx]
            # 递归最接近连通域的轮廓点, 找到最接近当前轮廓的位置, 然后画线
            near_point = None
            distance = np.inf
            for point_ in near_contour:
                point_ = point_[0]
                distance_ = np.sqrt(
                    np.sum(np.power(connect_point - point_, 2))
                )  # 计算两点距离
                if distance_ < distance:
                    near_point = point_
                    distance = distance_

            # 连线
            label_mask = cv2.line(
                label_mask,
                connect_point.astype(np.int32),
                near_point.astype(np.int32),
                255,
                2,
            )
            break  # 每次连接两个轮廓后label_mask都发生了改变, 因此需要break后重新计算轮廓和连通域

    return label_mask


def is_intersect(pointer_tail, pointer_head, points):
    """判断射线是否与曲线轮廓的两个端点形成的线段相交
    input:
        pointer_tail:指针尾部,
        pointer_head:指针头部,
        points:曲线轮廓二维点集
    return:
    """
    for i in range(len(points) - 1):
        line_start = points[i]
        line_end = points[i + 1]

        if is_intersect_segment(pointer_tail, pointer_head, line_start, line_end):
            return True

    return False


def is_intersect_segment(pointer_tail, pointer_head, seg_start, seg_end):
    """利用向量判断射线是否与曲线轮廓线段相交
    input:
        pointer_tail:指针尾部,
        pointer_head:指针头部,
        seg_start:曲线轮廓起点,
        seg_end:曲线轮廓终点，
    return:
    """
    # 计算曲面切点切线向量
    dir_segment = seg_end - seg_start
    # 表盘指针射线交点
    cross = abs(pointer_head[1] * dir_segment[0] - pointer_head[0] * dir_segment[1])
    # 平行或共线
    if cross == 0:
        return False, None
    t = (
        (seg_start[0] - pointer_tail[0]) * dir_segment[0]
        + (seg_start[1] - pointer_tail[1]) * dir_segment[1]
    ) / cross
    # 相交
    if t >= 0 and t <= 1:
        return True
    return False


def find_intersection(curve_points, pointer_head, pointer_tail):
    """求射线与曲线轮廓交点
    input:
        curve_points: 曲线轮廓二维点集
        pointer_head: 指针头部
        pointer_tail: 指针尾部
    return:
        射线与曲线轮廓的交点点集: intersections
    """

    dist = np.inf
    n = len(curve_points)
    # 获取指针中点坐标
    pointer_mid = (pointer_head + pointer_tail) / 2
    for i in range(n):
        # 获取当前线段的端点
        p1 = curve_points[i % n]
        p2 = curve_points[(i + 1) % n]
        p3 = curve_points[(i + 2) % n]
        p4 = curve_points[(i + 3) % n]
        p5 = curve_points[(i + 4) % n]
        p_ = (p1 + p2 + p3) / 3
        p_mid = (p1 + p2 + p3 + p4 + p5) / 5
        # 计算中点到射线的最小距离
        dist_ = np.sqrt(np.sum(np.power(pointer_mid - p_mid, 2)))
        # dist_ = np.sqrt(np.sum(np.power(pointer_head - p_mid, 2)))
        # 循环搜索指针中点到曲线轮廓上距离最小的坐标点
        if dist_ < dist:
            # 求交点坐标
            # intersection = p_mid
            intersection = pointer_mid
            # intersection = (p_mid + pointer_mid)/2
            # 求交点到点集的最小距离
            dist = dist_
            closest_index = i
    return intersection


def calculate_total_arclen(points):
    """插值计算曲线轮廓总长
    input:
        points: 曲线轮廓二维点集
        closest_index: 交点在曲线点集上的序号
    return:
        np.array(total_arclen): 曲线轮廓弧长
    """
    cumulative_lengths = [0]
    for i in range(1, len(points)):
        # 累加在points[i-1]与points[i]之间的弦长
        dx = points[i, 0] - points[i - 1, 0]
        dy = points[i, 1] - points[i - 1, 1]
        chord_length = np.sqrt(dx**2 + dy**2)
        # 计算曲线轮廓弧长
        cumulative_lengths.append(cumulative_lengths[-1] + chord_length)
    return np.array(cumulative_lengths)


def calculate_arclen_to_intersection(points, closest_index):
    """插值计算曲线轮廓总长
    input:
        points: 曲线轮廓二维点集
        closest_index: 交点在曲线点集上的序号
    return:
        np.array(total_arclen): 曲线轮廓弧长
    """
    arclen_to_intersection = [0]
    for i in range(1, int(closest_index)):  # int(closest_index)
        # 累加在points[i-1]与points[i]之间的弦长
        dx = points[i, 0] - points[i - 1, 0]
        dy = points[i, 1] - points[i - 1, 1]
        chord_length = np.sqrt(dx**2 + dy**2)
        # 计算曲线轮廓弧长
        arclen_to_intersection.append(arclen_to_intersection[-1] + chord_length)
    return np.array(arclen_to_intersection)


def find_arc_length_to_point(points, intersection):
    """插值法求点集上交点弧长百分比
    input:
        points: 曲线轮廓二维点集
        intersection: 表针射线与曲线的交点
    return:
        ratio_arclen: 起点到表针交点的弧长在总弧长上的所占百分比
    """

    # 累加计算总弧长
    total_arclen = calculate_total_arclen(points)
    # 找到交点在点集中最小值所在的像素下标
    distances = np.sqrt(np.sum((points - intersection) ** 2, axis=1))
    closest_index = np.argmin(distances)
    # 累加计算刻度弧长
    arclen_to_intersection = calculate_arclen_to_intersection(points, closest_index)
    # 计算弧长百分比
    ratio_arclen = max(arclen_to_intersection) / max(total_arclen)

    return ratio_arclen
